package com.ybox.gateway.filter;

import com.ybox.common.core.constant.FilterOrderConstant;
import com.ybox.gateway.context.GatewayContext;
import io.netty.buffer.ByteBufAllocator;
import org.apache.commons.compress.utils.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.ResolvableType;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpMethod;
import org.springframework.http.InvalidMediaTypeException;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.codec.multipart.Part;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.CollectionUtils;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author ych
 */
@Component
public class ApiRequestFilter implements GlobalFilter, Ordered {
    private static final Logger log = LoggerFactory.getLogger(ApiRequestFilter.class);

    private static AntPathMatcher antPathMatcher;

    static {
        antPathMatcher = new AntPathMatcher();
    }

    /**
     * default HttpMessageReader
     */
    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();

    private static final ResolvableType MULTIPART_DATA_TYPE = ResolvableType.forClassWithGenerics(MultiValueMap.class, String.class, Part.class);

    private static final Mono<MultiValueMap<String, Part>> EMPTY_MULTIPART_DATA = Mono.just(CollectionUtils.unmodifiableMultiValueMap(new LinkedMultiValueMap<String, Part>(0))).cache();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        String url = request.getURI().getPath();

        if (request.getMethod() == HttpMethod.GET || request.getMethod() == HttpMethod.DELETE) {
            // get请求 处理参数
            return handleGetMethod(exchange, chain, request);
        }

        if (request.getMethod() == HttpMethod.POST || request.getMethod() == HttpMethod.PUT) {
            // post put请求 处理参数
            return handlePostMethod(exchange, chain, request);
        }

        return chain.filter(exchange);
    }

    /**
     * get请求 处理参数
     *
     * @param exchange
     * @param chain
     * @param request
     * @return
     */
    private Mono<Void> handleGetMethod(ServerWebExchange exchange, GatewayFilterChain chain, ServerHttpRequest request) {
        // TODO 暂时不做处理
        GatewayContext gatewayContext = new GatewayContext();
        gatewayContext.setPath(request.getPath().pathWithinApplication().value());
        return chain.filter(exchange);
    }

    /**
     * post请求 校验参数
     *
     * @param exchange
     * @param chain
     * @param request
     * @return
     */
    private Mono<Void> handlePostMethod(ServerWebExchange exchange, GatewayFilterChain chain, ServerHttpRequest request) {
        GatewayContext gatewayContext = new GatewayContext();
        gatewayContext.setPath(request.getPath().pathWithinApplication().value());
        /**
         * save gateway context into exchange
         */
        exchange.getAttributes().put(GatewayContext.CACHE_GATEWAY_CONTEXT, gatewayContext);

        MediaType contentType = request.getHeaders().getContentType();
        if (MediaType.APPLICATION_JSON.equals(contentType)
                || MediaType.APPLICATION_JSON_UTF8.equals(contentType)) {
            // 请求内容为 application json

            // 重新构造 请求体
            return readJsonBody(exchange, chain, gatewayContext);
        }

        if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(contentType)) {
            // 请求内容为 form data
            return readFormData(exchange, chain, gatewayContext);
        }
        return chain.filter(exchange);
    }

    /**
     * post 请求
     * 重新构造 请求体
     *
     * @param exchange
     * @param chain
     * @param gatewayContext
     * @return
     */
    private Mono<Void> readJsonBody(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    /*
                     * read the body Flux<DataBuffer>, and release the buffer
                     * //TODO when SpringCloudGateway Version Release To G.SR2,this can be update with the new version's feature
                     * see PR https://github.com/spring-cloud/spring-cloud-gateway/pull/1095
                     */
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                        DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                        DataBufferUtils.retain(buffer);
                        return Mono.just(buffer);
                    });
                    // repackage ServerHttpRequest
                    ServerHttpRequest mutatedRequest =
                            new ServerHttpRequestDecorator(exchange.getRequest()) {
                                @Override
                                public Flux<DataBuffer> getBody() {
                                    return cachedFlux;
                                }
                            };
                    // mutate exchage with new ServerHttpRequest
                    ServerWebExchange mutatedExchange = exchange.mutate().request(mutatedRequest).build();
                    // read body string with default messageReaders
                    return ServerRequest.create(mutatedExchange, messageReaders)
                            .bodyToMono(String.class)
                            .doOnNext(objectValue -> {
                                // save body into gatewayContext
                                gatewayContext.setCacheBody(objectValue);
                            })
                            .then(chain.filter(mutatedExchange));
                });
    }

    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        return exchange.getRequest().getBody().collectList().flatMap(dataBuffers -> {
            final byte[] totalBytes = dataBuffers.stream().map(dataBuffer -> {
                try {
                    final byte[] bytes = IOUtils.toByteArray(dataBuffer.asInputStream());
                    //  System.out.println(new String(bytes));
                    return bytes;
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }).reduce(this::addBytes).get();
            final ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                @Override
                public Flux<DataBuffer> getBody() {
                    return Flux.just(buffer(totalBytes));
                }
            };
            final ServerCodecConfigurer configurer = ServerCodecConfigurer.create();
            final Mono<MultiValueMap<String, Part>> multiValueMapMono = repackageMultipartData(decorator, configurer);
            return multiValueMapMono.flatMap(part -> {
                for (String key : part.keySet()) {
                    // 如果为文件时 则进入下一次循环
                    if ("file".equals(key)) {
                        continue;
                    }
                    part.getFirst(key).content().subscribe(buffer -> {
                        final byte[] bytes = new byte[buffer.readableByteCount()];
                        buffer.read(bytes);
                        DataBufferUtils.release(buffer);
                        try {
                            final String bodyString = new String(bytes, "utf-8");
                            gatewayContext.setCacheBody(bodyString);
                        } catch (UnsupportedEncodingException e) {
                            e.printStackTrace();
                        }
                    });
                }
                return chain.filter(exchange.mutate().request(decorator).build());
            });
        });
    }

    @SuppressWarnings("unchecked")
    private static Mono<MultiValueMap<String, Part>> repackageMultipartData(ServerHttpRequest request, ServerCodecConfigurer configurer) {
        try {
            final MediaType contentType = request.getHeaders().getContentType();
            if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(contentType)) {
                return ((HttpMessageReader<MultiValueMap<String, Part>>) configurer.getReaders().stream().filter(reader -> reader.canRead(MULTIPART_DATA_TYPE, MediaType.MULTIPART_FORM_DATA))
                        .findFirst().orElseThrow(() -> new IllegalStateException("No multipart HttpMessageReader."))).readMono(MULTIPART_DATA_TYPE, request, Collections.emptyMap())
                        .switchIfEmpty(EMPTY_MULTIPART_DATA).cache();
            }
        } catch (InvalidMediaTypeException ex) {
            // Ignore
        }
        return EMPTY_MULTIPART_DATA;
    }

    /**
     * addBytes.
     *
     * @param first  first
     * @param second second
     * @return byte
     */
    public byte[] addBytes(byte[] first, byte[] second) {
        final byte[] result = Arrays.copyOf(first, first.length + second.length);
        System.arraycopy(second, 0, result, first.length, second.length);
        return result;
    }

    private DataBuffer buffer(byte[] bytes) {
        final NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
        final DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
        buffer.write(bytes);
        return buffer;
    }


    @Override
    public int getOrder() {
        return FilterOrderConstant.getOrder(this.getClass().getName());
    }
}

